Skip to content

PYTHON-5215 Add an asyncio.Protocol implementation for KMS #2460

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 37 commits into from
Aug 19, 2025
Merged

PYTHON-5215 Add an asyncio.Protocol implementation for KMS #2460

merged 37 commits into from
Aug 19, 2025

Conversation

blink1073
Copy link
Member

@blink1073 blink1073 commented Aug 6, 2025 •
edited
Loading

See benchmark gist.

Benchmark Results:

Before: 4.93s, 5.26s
After: 4.93s, 5.05s

Depends on mongodb-labs/drivers-evergreen-tools#679

@blink1073 blink1073 requested a review from NoahStapp August 6, 2025 19:30
@blink1073 blink1073 requested a review from a team as a code owner August 6, 2025 19:30
@blink1073
Copy link
Member Author

I'm debugging two failures:

test.asynchronous.test_connection_monitoring.AsyncTestCMAP.test_connection_monitoring_pool_clear_interrupting_pending_connections_clear_with_interruptInUseConnections___true_closes_pending_connections
test.asynchronous.test_connection_monitoring.AsyncTestCMAP.test_connection_monitoring_pool_create_min_size_error_error_during_minPoolSize_population_clears_pool

@blink1073 blink1073 marked this pull request as draft August 7, 2025 00:24
@blink1073 blink1073 removed the request for review from NoahStapp August 7, 2025 00:24
@blink1073
Copy link
Member Author

I realized the benchmark test wasn't actually triggering the protocol -> I'm tweaking things locally

NoahStapp
NoahStapp requested changes Aug 14, 2025
View reviewed changes
test/asynchronous/test_collection.py Outdated
@@ -335,6 +335,8 @@ async def test_create_index(self):
await db.test.create_index(["hello", ("world", DESCENDING)])
await db.test.create_index({"hello": 1}.items()) # type:ignore[arg-type]

# TODO: PYTHON-5491 - remove version max
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change should be in a separate PR.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it was, I just updated this branch.

pymongo/network_layer.py Outdated
# Reuse the active buffer if it has space.
if len(self._buffers):
buffer = self._buffers[-1]
if len(buffer.buffer) - buffer.end_index > sizehint:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we're setting sizehint to be at least 16384 always, is this check worth doing in the first place? I'd expect us to rarely reuse the active buffer since we'll usually have a buffer of size 16384 and a sizehint of 16384.

@blink1073
Copy link
Member Author

If we're setting sizehint to be at least 16384 always, is this check worth doing in the first place? I'd expect us to rarely reuse the active buffer since we'll usually have a buffer of size 16384 and a sizehint of 16384.

The actual sizehint in practice was on the order of the bytes being read from the buffer (typically less than 1000). Using the buffered protocol at all here is a bit of a mismatch imho.

@blink1073 blink1073 requested a review from NoahStapp August 14, 2025 19:00
@NoahStapp
Copy link
Contributor

If we're setting sizehint to be at least 16384 always, is this check worth doing in the first place? I'd expect us to rarely reuse the active buffer since we'll usually have a buffer of size 16384 and a sizehint of 16384.

The actual sizehint in practice was on the order of the bytes being read from the buffer (typically less than 1000). Using the buffered protocol at all here is a bit of a mismatch imho.

How long would refactoring to not use buffered take? No reason to use the lower-level API if we don't need to.

@blink1073
Copy link
Member Author

How long would refactoring to not use buffered take? No reason to use the lower-level API if we don't need to.

It's actually dead simple, I did it along the way when I was debugging a race condition.

@blink1073
Copy link
Member Author

It's actually dead simple, I did it along the way when I was debugging a race condition.

I'll push a commit in the morning for comparison, we can always revert.

@blink1073
Copy link
Member Author

I'm happy with the simplification. The tests are passing locally, this is ready for another look.

NoahStapp
NoahStapp reviewed Aug 15, 2025
View reviewed changes
Copy link
Contributor

@NoahStapp NoahStapp left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you schedule a full Evergreen run? We should ensure there's no regressions introduced here by accident.

Are the benchmark results for KMS significantly different between the two Protocol implementations?

@blink1073
Copy link
Member Author

Full patch build: https://spruce.mongodb.com/version/689f5483e112170007b0ce9f/tasks?sorts=STATUS%3AASC%3BBASE_STATUS%3ADESC

I updated the timings in the PR description, there no significant change.

@blink1073 blink1073 requested a review from NoahStapp August 15, 2025 15:39
@blink1073
Copy link
Member Author

Okay there is one legit bug in test.test_connection_logging.TestConnectionLoggingConnectionLogging.test_Connection_checkout_fails_due_to_error_establishing_connection. I'll defer looking at that until next week to focus on greener build tasks.

@blink1073
Copy link
Member Author

Here's a new patch build, failures are existing flakiness issues or tracked in PYTHON-5502.

https://spruce.mongodb.com/version/68a3434b33ee5e00072767c9/tasks?sorts=STATUS%3AASC%3BBASE_STATUS%3ADESC

NoahStapp
NoahStapp reviewed Aug 18, 2025
View reviewed changes
pymongo/synchronous/pool.py
@@ -124,7 +124,89 @@ def _set_non_inheritable_non_atomic(fd: int) -> None: # noqa: ARG001
_IS_SYNC = True


class Connection:
class BaseConnection:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't utilize any of the BaseConnection abstractions for the sync API, correct? The change is purely for compatibility with the async API?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No it is used by both sync and async KMS, it is what gets returned by the _connect_kms helper function.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I meant that we don't have a separate KMS networking interface for sync. Just wanted to confirm my understanding.

@blink1073 blink1073 requested a review from NoahStapp August 18, 2025 20:44
NoahStapp
NoahStapp previously approved these changes Aug 18, 2025
View reviewed changes
@blink1073 blink1073 requested a review from NoahStapp August 18, 2025 23:59
@blink1073
Copy link
Member Author

I had to fix merge conflicts

@blink1073 blink1073 merged commit e4b7eb5 into mongodb:master Aug 19, 2025
79 checks passed
@blink1073 blink1073 deleted the PYTHON-5215 branch August 19, 2025 13:45
sleepyStick pushed a commit to sleepyStick/mongo-python-driver that referenced this pull request Aug 19, 2025
@blink1073 @sleepyStick
PYTHON-5215 Add an asyncio.Protocol implementation for KMS (mongodb#2460
e33e393 
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@NoahStapp NoahStapp NoahStapp approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@blink1073 @NoahStapp